import os
import pandas
import pickle
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import FuncFormatter
from matplotlib.ticker import ScalarFormatter

def moving_average(rewards, window_size=10):
    """
    Calculate the moving average of rewards.

    Parameters:
    - rewards (list or np.array): Sequence of rewards to smooth.
    - window_size (int): Size of the moving window.

    Returns:
    - np.array: Smoothed rewards using moving average.
    """
    rewards = np.array(rewards)
    cumsum = np.cumsum(np.insert(rewards, 0, 0)) 
    moving_avg = (cumsum[window_size:] - cumsum[:-window_size]) / window_size
    return moving_avg

HOME = os.getcwd()
with open(f'{HOME}/TD3/TD3_record.pkl', 'rb') as f:
    td3_data = pickle.load(f)

with open(f'{HOME}/SAC/sac_50_record_add.pkl', 'rb') as f:
    sacn_data_50 = pickle.load(f)

with open(f'{HOME}/SAC/sac_20_record_add.pkl', 'rb') as f:
    sacn_data_20 = pickle.load(f)

with open(f'{HOME}/SAC/sac_10_record_add.pkl', 'rb') as f:
    sacn_data_10 = pickle.load(f)

with open(f'{HOME}/SAC/sac_5_record_add.pkl', 'rb') as f:
    sacn_data_5 = pickle.load(f)

smoothed_td3_data = moving_average(td3_data['score'][:200], window_size=15)
smoothed_sacn_data_50 = moving_average(sacn_data_50['score'][:200], window_size=15)
smoothed_sacn_data_20 = moving_average(sacn_data_20['score'][:200], window_size=15)
smoothed_sacn_data_10 = moving_average(sacn_data_10['score'][:200], window_size=15)
smoothed_sacn_data_5 = moving_average(sacn_data_5['score'][:200], window_size=15)


plt.rcParams['font.size']=16
fig, axe = plt.subplots(figsize=(8,6))

axe.plot(smoothed_td3_data, label='TD3')
axe.plot(smoothed_sacn_data_50, label='SAC-50')
axe.plot(smoothed_sacn_data_20, label='SAC-20')
axe.plot(smoothed_sacn_data_10, label='SAC-10')
# axe.plot(smoothed_sacn_data_5, label='SAC-5')
axe.set_xlabel('Epochs')
axe.set_ylabel('Score')
axe.legend()

# 设置 y 轴刻度格式化
formatter = ScalarFormatter(useMathText=True)  # 使用科学计数法显示
formatter.set_powerlimits((-6, 6))  # 控制何时使用科学计数法
axe.yaxis.set_major_formatter(formatter)

# 强制显示 offset text
axe.ticklabel_format(axis="y", style="scientific")

# 将 offset text 移动到 figure 的左上角
axe.yaxis.get_offset_text().set_position((0, 1.05))  # 偏移设置
axe.yaxis.get_offset_text().set_horizontalalignment('left')  # 左对齐
# 显示图表
plt.tight_layout()
# formatter = FuncFormatter(scientific_notation_formatter)
# axe.yaxis.set_major_formatter(formatter)
plt.savefig(f'{HOME}/data_deal/plot_score.eps', bbox_inches='tight', dpi=300)
plt.show()